function [outSt,logL]=kfilterNanSplitsGeneral(parvec,model,dataStru,flags,filterStru)
%function kfilterNanSplitsGeneral(parvec,model,dataStru,flags,filterStru)
% ====================================================================

%% Relative to kfilterNanSplitsGeneral; 

%% Replacements 
%% model.handle         replaces funcmod
%% dataStru.data        replaces dataStru.data
%% dataStru.trainVec    replaces trainVec 
%% model.addsol         replaces addsol 
%% model.solveOptions   replaces solveopt 


%% Additional Inputs 
%% filterStru           .tauVec 
%%                      .aZero          [ns 1]   state from previous run
%%                      .pZero          [ns ns]  initial covariance from previous run  
%%                      .aZeroDecomp    [ns nx]  initial decomposition from previous run 
%%
%% flags                .initialCond  if Initial Condition provided 
%% 
%% outSt
% outSt.sttNoInitial=stt        [nobs ns]    vector of one-sided states
% outSt.pttNoInitial=ptt        [ns ns nobs] vector of RMSE for stt 
% outSt.filteredSt=stt          [nobs ns]    vector of one-sided states
% outSt.aLast=stt(end,:)'       [ns    1]    last state s_{T}|T
% outSt.pLast=ptt(:,:,end)      [ns   ns]    last rmse  P_{T}|T 
% outSt.smoothSt=smooth_st      [nobs ns]    vector of two-sides states
% outSt.innovations=etamat      [nobs nx]    vector of innovations 
% outSt.countSt=countStates     [nobs ns ns] vector of states decompositions
%                               1 page per shock 
% outSt.countObs=countObs       [nobs ny]    vector of observables
%                               decompositions, 1 page per shock 
% outSt.decompInitialSt=sOneSmoothPerShock [ns nx] decomposition of initial
%                               state 
% outSt.forecastObs=yfor        [nobs ny] 1 step ahead forecast for observables  
% outSt.aZero=a0;               [ns    1] initial state 

% This is the general version of *kfilterNSplits.m* 
% 
%% *Note* 
%  Decomposition will NOT be exact, when intial state is provided as 
%  it will require passing also the decomposition of that state and it's 
%  associated variance covariance
% 
%  A. Justiniano Feb 25 2014 
% ====================================================================
%% 1. Model solution
% Matrices of second sample stored in structure second 
[G,R,C,eu,SDX,Z,structOne,ssvec,structTwo]=...
    feval(model.handle,parvec,model.solveOptions,model.addsol);
if isequal(eu,[1;1])==0 
    error('Model is not determinate') 
end 
[T,dim.var]=size(dataStru.data);
dataStru.data=dataStru.data';
tauVec=model.addsol.tauVec; 
if length(tauVec)~=T 
    warning('model.addsol is Longer than T'); 
    quer('c'); 
end 
%% 3. Forward filter 
% 3.1. Dimensions and storage 
ns = size(G,1); 
nx = size(SDX,1); 

vt=zeros(dim.var,T);
finvt=zeros(dim.var,dim.var,T);
kpartg=zeros(ns,dim.var,T);
logLnc=zeros(T,1); 
yfor =zeros(dim.var,T);
% Matrices with one additional entrdataStru.data (initialization)
% to recover observables 
stt=zeros(ns,T+1);
ptt=zeros(ns,ns,T+1);
W          = eye(dim.var);
Zdim       = zeros(T,1); 
mat_obspos = zeros(T,dim.var);

% 3.2 Initialization 
if flags.initialCond==0
    stt(:,1) = zeros(ns,1);
    P0 = lyapunov_symm(G(:,:,tauVec(1)),...
        R(:,:,tauVec(1))*(SDX(:,:,tauVec(1))')...
        *SDX(:,:,tauVec(1))*(R(:,:,tauVec(1))'));
    ptt(:,:,1)=P0;
else
    P0=filterStru.pZero;
    stt(:,1)=filterStru.aZero; 
    ptt(:,:,1)=P0; 
end 

% 3.3 Start Forward Filter using KF_DK 
for ii=1:T;
    
    % Handling of missing observations
    ytt=dataStru.data(:,ii);
    
    % Determine W and position of the NAN
    ind =~isnan(ytt);
    rowt =find(~isnan(ytt));
    
    
    ytt=ytt(ind);
    Zdim(ii)=length(ytt);
    Ztt=W((ind==1),:)*Z(:,:,tauVec(ii));
    mat_obspos(ii,1:Zdim(ii))=rowt;
    dimt=( 1:Zdim(ii) );
    %% Demeaning is done here 
    ytt = ytt - Ztt*C(:,tauVec(ii));
    
%     if addsol.tauVec(ii)==2 
%         disp('Here'); 
%     end 
    
    yfor(dimt,ii) =Ztt*(G(:,:,tauVec(ii))*stt(:,ii) + C(:,tauVec(ii)));
    
    [stt(:,ii+1),ptt(:,:,ii+1),logLnc(ii),vt(dimt,ii),finvt(dimt,dimt,ii),...
        kpartg(:,dimt,ii)]=feval(@kf_dk,ytt,Ztt,...%Z(:,:,tauVec(ii)),...
        stt(:,ii),ptt(:,:,ii),G(:,:,tauVec(ii)),...
        R(:,:,tauVec(ii))*(SDX(:,:,tauVec(ii))'));
end

%% 4. Likelihood with Integration Constant  
outSt.logLncFull=logLnc; 
logLnc=logLnc(dataStru.trainVec(1):dataStru.trainVec(2)); 
logL=-0.5*sum(Zdim)*log(2*pi)+sum(logLnc);

%% 5. Truncate filters and obtain initial observations 
yferr=vt';
yfor =yfor'; 
stt=stt(:,2:end); ptt=ptt(:,:,2:end); 

%% 6. Disturbance smoother with TV matrices 
% Obtain the Innovations using a disturbance smoother 
etamat=zeros(nx,T); 
smooth_st=zeros(ns,T); 
rmat=zeros(ns,T); 

%% 6.1 Initialize RSTAR & start at t=Nobs 
rstar=zeros(ns,1); 
[rstar,etamat(:,end)]=smoothdis(rstar,...
    (SDX(:,:,tauVec(end))')*SDX(:,:,tauVec(end)),R(:,:,tauVec(end))',...
    Ztt',...%Z(dimt,:,tauVec(end))',...
    finvt(dimt,dimt,end),zeros(ns),vt(dimt,end));
smooth_st(:,end)=stt(:,end)+ptt(:,:,end)*rstar; 
rmat(:,end)=rstar; 

%% 6.2 Begin Backward recursion 
for ii=(T-1):-1:1;
    
    % Varying dimension in the backward recurions
    dimt=( 1:Zdim(ii) );
    % [Ztt]'= [Wtt*Z]' = = Z'*Wtt'
    Ztt = (Z(:,:,tauVec(ii))')*( W( mat_obspos(ii,1:Zdim(ii)),:)');
    
    [rstar,etamat(:,ii)]=smoothdis(rstar,...
        (SDX(:,:,tauVec(ii))')*SDX(:,:,tauVec(ii)),...
        R(:,:,tauVec(ii))',Ztt,finvt(dimt,dimt,ii),...
        ((G(:,:,tauVec(ii+1))-G(:,:,tauVec(ii+1))*...
        kpartg(:,dimt,ii)*Ztt')'),vt(dimt,ii));
    smooth_st(:,ii)=stt(:,ii)+ptt(:,:,ii)*rstar;
    rmat(:,ii)=rstar; 
end
a0 = P0*G(:,:,tauVec(1))'*rstar;   % Note: this is only correct in the case where a0 = zeros(ns,1) 
etamat=(etamat)';
smooth_st=(smooth_st)'; 
stt  =stt'; 


%% 7. Check Smoother 
% Check that Smooth States are identical if using disturbance smoother (above) vs. state smoother (below) 
% and if can also recover the observables
tol=1e-5; 
% State at time zero give by GG1*azero+Pzero*rstar;
[yCheck,smoothCheck]=kfilterRegSplitSimulation(a0,etamat'); 
disp('Max Discrepancy Smooth vs. Actual Data'); 
Ydem=zeros(dim.var,T); 
for ii=1:T 
    Ydem(:,ii)=dataStru.data(:,ii)- Z(:,:,tauVec(ii))*C(:,tauVec(ii));
end 
maxdifY=comparemat(yCheck,Ydem'); 
disp('Max Discrepancy State and Innovation Smoother'); 
maxdifS=comparemat(smoothCheck,smooth_st); 
if maxdifY > tol || maxdifS > tol 
    warning('Smoother discrepancy exceeds tolerance') 
end 

%% 8. Initial Variance and State per shock (used below for the counterfactual decompositions)
% vInitialPershock: Initial Variance, decomposed per shock 
% sOnePerShock, State at time 1 decomposed, decomposed per shock
vInitialPerShock=zeros(ns,ns,nx);
sZeroSmoothPerShock=zeros(ns,nx); 
sOneSmoothPerShock =zeros(ns,nx);
for ii=1:nx
    
    etatemp = zeros(nx,1);
    etatemp(ii) = etamat(1,ii);
    
    if flags.initialCond==0
        vInitialPerShock(:,:,ii)=lyapunov_symm(G(:,:,tauVec(1)),...
            R(:,ii,tauVec(1))*SDX(ii,ii,tauVec(1))*...
            SDX(ii,ii,tauVec(1))*(R(:,ii,tauVec(1))'));
        sZeroSmoothPerShock(:,ii) = vInitialPerShock(:,:,ii)*G(:,:,tauVec(1))'*rstar;
      
        sOneSmoothPerShock(:,ii)=G(:,:,tauVec(1))*sZeroSmoothPerShock(:,ii)...
            +R(:,:,tauVec(1))*etatemp;
        
    else
        sZeroSmoothPerShock(:,ii)=filterStru.aZeroDecomp(:,ii);
        sOneSmoothPerShock(:,ii)=G(:,:,tauVec(1))*sZeroSmoothPerShock(:,ii)...
            +R(:,:,tauVec(1))*etatemp;
    end
    
end
% This will work for a0 even if flags.initialCond==1
maxdifSzero=comparemat(sum(sZeroSmoothPerShock,2),a0);
% Not sure this will work if flags.initialCond==1
maxdifSone =comparemat(sum(sOneSmoothPerShock,2),smooth_st(1,:)' );
if maxdifSzero > tol; warning('Cannot decompose sZeroPerShock'); end
if maxdifSone > tol; warning('Cannot decompose sOnePerShock'); end

%% 9. Counterfactual Decomposition 
% Generate states by feeding each shock at a time, ensure that it recovers
% the original state 
disp('Begin Counterfactual')
countStates=zeros(T,ns,nx);
countObs   =zeros(T,dim.var,nx);
for ii=1:nx 
    etaTemp=zeros(T,nx); 
    etaTemp(:,ii)=etamat(:,ii); 
    [countObs(:,:,ii),countStates(:,:,ii)]=kfilterRegSplitSimulation(sZeroSmoothPerShock(:,ii),etaTemp');
end 
disp('Maximum Discrepancy Counterfactual States and Smooth States') 
maxdifCount=comparemat(sum(countStates,3),smooth_st); 
if maxdifCount > tol; warning('Counterfactual States do not recover smooth states'); end 

% Use this to debug testlikelInitial.m 
outSt.sttNoInitial=stt; 
outSt.pttNoInitial=ptt; 
outSt.filteredSt=stt; 
outSt.aLast=stt(end,:)'; 
outSt.pLast=ptt(:,:,end); 

outSt.smoothSt=smooth_st; 
outSt.innovations=etamat; 
outSt.countSt=countStates; 
outSt.countObs=countObs; 
outSt.decompInitialSt=sOneSmoothPerShock; 
outSt.forecastObs=yfor; 
outSt.aZero=a0; 


%% Subroutine kfilterRegSplitSimulation allows to simulate the model  
% Inputs 
% sInitial: [ns 1] Initial State 
% innovMat: [nx T] matrix of innovations 
% By being a sub-routine it has access to all variables defined above 
% Be-careful not to use an index (ii,jj) used above or to repeat variable
% names
    function [ySim,sSim]=kfilterRegSplitSimulation(sInitial,innovMat) 
        if ~isequal(size(innovMat),[nx T]) 
            error('Input innovMat must be [nx T]') 
        end 
        sSim=zeros(ns,T);
        ySim=zeros(size(dataStru.data));
        sSim(:,1)=G(:,:,tauVec(1))*sInitial + ...
            R(:,:,tauVec(1))*innovMat(:,1);
        ySim(:,1)=Z(:,:,tauVec(1))*sSim(:,1); %+C(:,tauVec(1)));
        for kk=2:T;
            sSim(:,kk)=G(:,:,tauVec(kk))*sSim(:,kk-1)+...
                R(:,:,tauVec(kk))*innovMat(:,kk);
            ySim(:,kk)=Z(:,:,tauVec(kk))*sSim(:,kk); %+C(:,tauVec(kk)));
        end
        ySim=ySim'; 
        sSim=sSim'; 
    end
%% End of File
end 